# -*- coding: utf-8 -*-  
'''
字典

Created on 2021年10月1日
@author: luoyi
'''
import pickle
import threading

import utils.conf as conf


#    字典
class Dictionaries:
    _INSTANCE_LOCK = threading.Lock()
    
    @classmethod
    def instance(cls, *args, **kwargs):
        if (not hasattr(Dictionaries, '_instance')):
            with Dictionaries._INSTANCE_LOCK:
                if (not hasattr(Dictionaries, '_instance')):
                    Dictionaries._instance = Dictionaries(*args, **kwargs)
                    pass
                pass
            pass
        return Dictionaries._instance
    
    
    def __init__(self, 
                 dictionaries_chchars_path=conf.DICTIONARY.get_dictionaries_chchars_path(),
                 dictionaries_word2dict_pkl_path=conf.DICTIONARY.get_dictionaries_word2dict_pkl_path(),
                 dictionaries_dict2word_pkl_path=conf.DICTIONARY.get_dictionaries_dict2word_pkl_path()):
        self._dictionaries_chchars_path = dictionaries_chchars_path
        self._dictionaries_word2dict_pkl_path = dictionaries_word2dict_pkl_path
        self._dictionaries_dict2word_pkl_path = dictionaries_dict2word_pkl_path
        
        #    当词取不到id时，是否追加新id
        self._is_auto_append = False
        
        self._w2id = {}
        self._id2w = {}
        pass
    
    
    #    写入pkl文件
    def write_pkl(self):
        #    字典中的默认值
        d_w2id = {'[PAD]':0, '[CLS]':1, '[SEP]':2, '[UNK]':3, '[MASK]':4, '[EOS]':5}
        d_id2w = {v:k  for k, v in d_w2id.items()}
        
        #    从原始文件中加载字典
        da_w2id, da_id2w = self.load_from_original(reserved=len(d_w2id))
        d_w2id.update(da_w2id)
        d_id2w.update(da_id2w)
        
        #    写入磁盘
        with open(file=self._dictionaries_word2dict_pkl_path, mode='wb') as fw: pickle.dump(d_w2id, fw)
        with open(file=self._dictionaries_dict2word_pkl_path, mode='wb') as fw: pickle.dump(d_id2w, fw)
        pass
    #    从原始文件中加载字典
    def load_from_original(self, reserved=5, max_empty=10):
        fr = open(file=self._dictionaries_chchars_path, mode='r', encoding='utf-8')
        i = reserved
        d_w2id = {}
        d_id2w = {}
        empty = 0
        while True:
            if (empty >= max_empty):
                break
            
            l = fr.readline()
            if (not l): 
                empty += 1
                continue
            if (len(l) == 0): 
                empty += 1
                continue
            
            empty = 0
            word = l.split(':')[0]
            if (d_w2id.get(word) is not None):
                continue
            
            d_w2id[word] = i
            d_id2w[i] = word
            i += 1
            pass
        fr.close()
        return d_w2id, d_id2w
    
    
    #    从pkl中加载字典
    def load_from_pkl(self):
        #    加载pkl文件
        with open(file=self._dictionaries_word2dict_pkl_path, mode='rb') as fr: self._w2id = pickle.load(fr)
        with open(file=self._dictionaries_dict2word_pkl_path, mode='rb') as fr: self._id2w = pickle.load(fr)
        pass
    
    
    #    取词的index（如果word_to_index里有这个词，则直接返回索引，否则扩充字典返回新的索引）
    def word2idx(self, word):
        idx = self._w2id.get(word)
        #    如果有就直接返回
        if (idx is not None): return idx
        
        #    如果不允许扩充字典，直接返回[UNK]
        if (not self._is_auto_append): 
            return self._w2id['[UNK]']
        
        #    否则扩充字典
        dict_size = len(self._w2id)
        self._w2id[word] = dict_size
        self._id2w[dict_size] = word
        
        return dict_size
    #    取词的index（如果word_to_index里有这个词，则直接返回索引，否则扩充字典返回新的索引）
    def word2idx_slist(self,word_list):
        idx_list = [self.word2idx(word) for word in word_list]
        return idx_list
    #    字典长度
    def size(self):
        return len(self._w2id)
    
    #    根据词的index取词
    def id2word(self, word_id):
        return self._id2w.get(word_id)
    #    根据词的index取词
    def id2word_idlist(self, word_ids):
        id_list = [self._id2w.get(word_id) for word_id in word_ids]
        return id_list
    pass

